/*
 * Copyright (C) 2009-2011 Big Muscle, http://strongdc.sf.net
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
 */

#include "StdAfx.h"
#include "SearchManager.h"

#include "Constants.h"
#include "DHT.h"
#include "IndexManager.h"
#include "Utils.h"

#include "../client/ClientManager.h"
#include "../client/SearchManager.h"
#include "../client/SearchResult.h"
#include "../client/SimpleXml.h"

namespace dht
{

Search::~Search()
{
	switch (type)
	{
		case TYPE_NODE:
			IndexManager::getInstance()->setPublish(true);
			break;
		case TYPE_STOREFILE:
			IndexManager::getInstance()->decPublishing();
			break;
	}
}

/*
 * Process this search request
 */
void Search::process()
{
	if (stopping)
		return;
		
	// no node to search
	if (possibleNodes.empty()/* || respondedNodes.size() >= MAX_SEARCH_RESULTS*/)
	{
		stopping = true;
		lifeTime = GET_TICK() + SEARCH_STOPTIME; // wait before deleting not to lose so much delayed results
		return;
	}
	
	// send search request to the first ALPHA closest nodes
	size_t nodesCount = min((size_t)SEARCH_ALPHA, possibleNodes.size());
	Node::Map::iterator it;
	for (size_t i = 0; i < nodesCount; i++)
	{
		it = possibleNodes.begin();
		Node::Ptr node = it->second;
		
		// move to tried and delete from possibles
		triedNodes[it->first] = node;
		possibleNodes.erase(it);
		
		// send SCH command
		AdcCommand cmd(AdcCommand::CMD_SCH, AdcCommand::TYPE_UDP);
		cmd.addParam("TR", term);
		cmd.addParam("TY", Util::toString(type));
		cmd.addParam("TO", token);
		
		//node->setTimeout();
		DHT::getInstance()->send(cmd, node->getIdentity().getIp(), static_cast<uint16_t>(Util::toInt(node->getIdentity().getUdpPort())), node->getUser()->getCID(), node->getUdpKey());
	}
}

SearchManager::SearchManager(void) : lastSearchFile(0)
{
}

SearchManager::~SearchManager(void)
{
}

/*
 * Performs node lookup in the network
 */
void SearchManager::findNode(const CID& cid)
{
	if (isAlreadySearchingFor(cid.toBase32()))
		return;
		
	Search* s = new Search();
	s->type = Search::TYPE_NODE;
	s->term = cid.toBase32();
	s->token = Util::toString(Util::rand());
	
	search(*s);
}

/*
 * Performs value lookup in the network
 */
void SearchManager::findFile(const string& tth, const string& token)
{
	// temporary fix to prevent UDP flood (search queue would be better here)
	if (GET_TICK() - lastSearchFile < 10000)
		return;
		
	if (isAlreadySearchingFor(tth))
		return;
		
	// do I have requested TTH in my store?
	//IndexManager::SourceList sources;
	//if(IndexManager::getInstance()->findResult(TTHValue(tth), sources))
	//{
	//  for(IndexManager::SourceList::const_iterator i = sources.begin(); i != sources.end(); ++i)
	//  {
	//      // create user as offline (only TCP connected users will be online)
	//      UserPtr u = ClientManager::getInstance()->getUser(i->getCID());
	//      u->setFlag(User::DHT);
	//
	//      // contact node that we are online and we want his info
	//      DHT::getInstance()->info(i->getIp(), i->getUdpPort(), true);
	//
	//      SearchResultPtr sr(new SearchResult(u, SearchResult::TYPE_FILE, 0, 0, i->getSize(), tth, "DHT", Util::emptyString, i->getIp(), TTHValue(tth), token));
	//      dcpp::SearchManager::getInstance()->fire(SearchManagerListener::SR(), sr);
	//  }
	//
	//  return;
	//}
	
	Search* s = new Search();
	s->type = Search::TYPE_FILE;
	s->term = tth;
	s->token = token;
	
	search(*s);
	
	lastSearchFile = GET_TICK();
}

/*
 * Performs node lookup to store key/value pair in the network
 */
void SearchManager::findStore(const string& tth, int64_t size, bool partial)
{
	if (isAlreadySearchingFor(tth))
	{
		IndexManager::getInstance()->decPublishing();
		return;
	}
	
	Search* s = new Search();
	s->type = Search::TYPE_STOREFILE;
	s->term = tth;
	s->filesize = size;
	s->partial = partial;
	s->token = Util::toString(Util::rand());
	
	search(*s);
}

/*
 * Performs general search operation in the network
 */
void SearchManager::search(Search& s)
{
	// set search lifetime
	s.lifeTime = GET_TICK();
	switch (s.type)
	{
		case Search::TYPE_FILE:
			s.lifeTime += SEARCHFILE_LIFETIME;
			break;
		case Search::TYPE_NODE:
			s.lifeTime += SEARCHNODE_LIFETIME;
			break;
		case Search::TYPE_STOREFILE:
			s.lifeTime += SEARCHSTOREFILE_LIFETIME;
			break;
	}
	
	// get nodes closest to requested ID
	DHT::getInstance()->getClosestNodes(CID(s.term), s.possibleNodes, 50, 3);
	
	if (s.possibleNodes.empty())
	{
		delete &s;
		return;
	}
	
	Lock l(cs);
	// store search
	searches[&s.token] = &s;
	
	s.process();
}

/*
 * Process incoming search request
 */
void SearchManager::processSearchRequest(const Node::Ptr& node, const AdcCommand& cmd)
{
	string token;
	if (!cmd.getParam("TO", 1, token))
		return; // missing search token?
		
	string term;
	if (!cmd.getParam("TR", 1, term))
		return; // nothing to search?
		
	string type;
	if (!cmd.getParam("TY", 1, type))
		return; // type not specified?
		
	AdcCommand res(AdcCommand::CMD_RES, AdcCommand::TYPE_UDP);
	res.addParam("TO", token);
	
	SimpleXML xml;
	xml.addTag("Nodes");
	xml.stepIn();
	
	bool empty = true;
	unsigned int searchType = Util::toInt(type);
	switch (searchType)
	{
		case Search::TYPE_FILE:
		{
			// check file hash in our database
			// if it's there, then select sources else do the same as node search
			IndexManager::SourceList sources;
			if (IndexManager::getInstance()->findResult(TTHValue(term), sources))
			{
				// yes, we got sources for this file
				unsigned int n = MAX_SEARCH_RESULTS;
				for (IndexManager::SourceList::const_iterator i = sources.begin(); i != sources.end() && n > 0; ++i, n--)
				{
					xml.addTag("Source");
					xml.addChildAttrib("CID", i->getCID().toBase32());
					xml.addChildAttrib("I4", i->getIp());
					xml.addChildAttrib("U4", i->getUdpPort());
					xml.addChildAttrib("SI", i->getSize());
					xml.addChildAttrib("PF", i->getPartial());
					
					empty = false;
				}
				break;
			}
		}
		default:
		{
			// maximum nodes in response is based on search type
			unsigned int count;
			switch (searchType)
			{
				case Search::TYPE_FILE:
					count = 2;
					break;
				case Search::TYPE_NODE:
					count = 10;
					break;
				case Search::TYPE_STOREFILE:
					count = 4;
					break;
				default:
					return; // unknown type
			}
			
			// get nodes closest to requested ID
			Node::Map nodes;
			DHT::getInstance()->getClosestNodes(CID(term), nodes, count, 2);
			
			// add nodelist in XML format
			for (Node::Map::const_iterator i = nodes.begin(); i != nodes.end(); ++i)
			{
				xml.addTag("Node");
				xml.addChildAttrib("CID", i->second->getUser()->getCID().toBase32());
				xml.addChildAttrib("I4", i->second->getIdentity().getIp());
				xml.addChildAttrib("U4", i->second->getIdentity().getUdpPort());
				
				empty = false;
			}
			
			break;
		}
	}
	
	xml.stepOut();
	
	if (empty)
		return; // no requested nodes found, don't send empty list
		
	string nodes;
	StringOutputStream sos(nodes);
	//sos.write(SimpleXML::utf8Header); // don't write header to save some bytes
	xml.toXML(&sos);
	
	res.addParam("NX", Utils::compressXML(nodes));
	
	// send search result
	DHT::getInstance()->send(res, node->getIdentity().getIp(), static_cast<uint16_t>(Util::toInt(node->getIdentity().getUdpPort())), node->getUser()->getCID(), node->getUdpKey());
}

/*
 * Process incoming search result
 */
void SearchManager::processSearchResult(const Node::Ptr& node, const AdcCommand& cmd)
{
	string token;
	if (!cmd.getParam("TO", 1, token))
		return; // missing search token?
		
	string nodes;
	if (!cmd.getParam("NX", 1, nodes))
		return; // missing search token?
		
	Lock l(cs);
	SearchMap::iterator i = searches.find(&token);
	if (i == searches.end())
	{
		// we didn't search for this
		return;
	}
	
	Search* s = i->second;
	
	// store this node
	s->respondedNodes.insert(std::make_pair(Utils::getDistance(node->getUser()->getCID(), CID(s->term)), node));
	
	try
	{
		SimpleXML xml;
		xml.fromXML(nodes);
		xml.stepIn();
		
		if (s->type == Search::TYPE_FILE) // this is response to TYPE_FILE, check sources first
		{
			// extract file sources
			while (xml.findChild("Source"))
			{
				const CID cid       = CID(xml.getChildAttrib("CID"));
				const string& i4    = xml.getChildAttrib("I4");
				uint16_t u4         = static_cast<uint16_t>(xml.getIntChildAttrib("U4"));
				int64_t size        = xml.getLongLongChildAttrib("SI");
				bool partial        = xml.getBoolChildAttrib("PF");
				
				// don't bother with invalid sources and private IPs
				if (cid.isZero() || ClientManager::getInstance()->getMe()->getCID() == cid || !Utils::isGoodIPPort(i4, u4))
					continue;
					
				// create user as offline (only TCP connected users will be online)
				Node::Ptr source = DHT::getInstance()->createNode(cid, i4, u4, false, false);
				
				if (partial)
				{
					if (!source->isOnline())
					{
						// node is not online, try to contact him
						DHT::getInstance()->info(i4, u4, DHT::PING | DHT::MAKE_ONLINE, cid, source->getUdpKey());
					}
					
					// ask for partial file
					AdcCommand cmd(AdcCommand::CMD_PSR, AdcCommand::TYPE_UDP);
					cmd.addParam("U4", Util::toString(::SearchManager::getInstance()->getPort()));
					cmd.addParam("TR", s->term);
					
					DHT::getInstance()->send(cmd, i4, u4, cid, source->getUdpKey());
				}
				else
				{
					// create search result: hub name+ip => "DHT", file name => TTH
					SearchResultPtr sr(new SearchResult(source->getUser(), SearchResult::TYPE_FILE, 0, 0, size, s->term, DHT::getInstance()->getHubName(), DHT::getInstance()->getHubUrl(), i4, TTHValue(s->term), token));
					if (!source->isOnline())
					{
						// node is not online, try to contact him if we didn't contact him recently
						if (searchResults.find(source->getUser()->getCID()) != searchResults.end())
							DHT::getInstance()->info(i4, u4, DHT::PING | DHT::MAKE_ONLINE, cid, source->getUdpKey());
							
						searchResults.insert(std::make_pair(source->getUser()->getCID(), std::make_pair(GET_TICK(), sr)));
					}
					else
					{
						sr->setSlots(Util::toInt(source->getIdentity().get("SL")));
						::SearchManager::getInstance()->fire(::SearchManagerListener::SR(), sr);
					}
				}
			}
			
			xml.resetCurrentChild();
		}
		
		// extract possible nodes
		unsigned int n = K;
		while (xml.findChild("Node") && n-- > 0)
		{
			CID cid = CID(xml.getChildAttrib("CID"));
			CID distance = Utils::getDistance(cid, CID(s->term));
			
			// don't bother with myself and nodes we've already tried or queued
			if (ClientManager::getInstance()->getMe()->getCID() == cid ||
			        s->possibleNodes.find(distance) != s->possibleNodes.end() ||
			        s->triedNodes.find(distance) != s->triedNodes.end())
			{
				continue;
			}
			
			const string& i4 = xml.getChildAttrib("I4");
			uint16_t u4 = static_cast<uint16_t>(xml.getIntChildAttrib("U4"));
			
			// don't bother with private IPs
			if (!Utils::isGoodIPPort(i4, u4))
				continue;
				
			// create unverified node
			// if this node already exists in our routing table, don't update it's ip/port for security reasons
			Node::Ptr node = DHT::getInstance()->createNode(cid, i4, u4, false, false);
			
			// node won't be accept for several reasons (invalid IP etc.)
			// if routing table is full, node can be accept
			bool isAcceptable = DHT::getInstance()->addNode(node, false);
			if (isAcceptable)
			{
				// update our list of possible nodes
				s->possibleNodes[distance] = node;
			}
		}
		
		xml.stepOut();
	}
	catch (const SimpleXMLException&)
	{
		// malformed node list
	}
}

/*
 * Sends publishing request
 */
void SearchManager::publishFile(const Node::Map& nodes, const string& tth, int64_t size, bool partial)
{
	// send PUB command to K nodes
	int n = K;
	for (Node::Map::const_iterator i = nodes.begin(); i != nodes.end() && n > 0; ++i, n--)
	{
		const Node::Ptr& node = i->second;
		
		AdcCommand cmd(AdcCommand::CMD_PUB, AdcCommand::TYPE_UDP);
		cmd.addParam("TR", tth);
		cmd.addParam("SI", Util::toString(size));
		
		if (partial)
			cmd.addParam("PF", "1");
			
		//i->second->setTimeout();
		DHT::getInstance()->send(cmd, node->getIdentity().getIp(), static_cast<uint16_t>(Util::toInt(node->getIdentity().getUdpPort())), node->getUser()->getCID(), node->getUdpKey());
	}
}

/*
 * Processes all running searches and removes long-time ones
 */
void SearchManager::processSearches()
{
	Lock l(cs);
	
	SearchMap::iterator it = searches.begin();
	while (it != searches.end())
	{
		Search* s = it->second;
		
		// process active search
		s->process();
		
		// remove long search
		if (s->lifeTime < GET_TICK())
		{
			// search timed out, stop it
			searches.erase(it++);
			
			if (s->type == Search::TYPE_STOREFILE)
			{
				publishFile(s->respondedNodes, s->term, s->filesize, s->partial);
			}
			
			delete s;
		}
		else
		{
			++it;
		}
	}
}

/*
 * Processes incoming search results
 */
bool SearchManager::processSearchResults(const UserPtr& user, size_t slots)
{
	bool ok = false;
	uint64_t tick = GET_TICK();
	
	ResultsMap::iterator it = searchResults.begin();
	while (it != searchResults.end())
	{
		if (it->first == user->getCID())
		{
			// user is online, process his result
			SearchResultPtr sr = it->second.second;
			sr->setSlots(slots); // slot count should be known now
			
			::SearchManager::getInstance()->fire(::SearchManagerListener::SR(), sr);
			searchResults.erase(it++);
			
			ok = true;
		}
		else if (it->second.first + 60 * 1000 <= tick)
		{
			// delete result from possibly offline users
			searchResults.erase(it++);
		}
		else
		{
			++it;
		}
	}
	
	return ok;
}

/*
 * Checks whether we are alreading searching for a term
 */
bool SearchManager::isAlreadySearchingFor(const string& term)
{
	Lock l(cs);
	for (SearchMap::const_iterator i = searches.begin(); i != searches.end(); ++i)
	{
		if (i->second->term == term)
			return true;
	}
	
	return false;
}

}